#include <vector>
#include <driver_types.h>
#include <device_launch_parameters.h>

#include "caffe/layers/eltwise_layer.hpp"

namespace caffe {

template <typename Ftype>
__global__ void MaxForward(const int nthreads, const Ftype* bottom_data_a,
    const Ftype* bottom_data_b, const int blob_idx, Ftype* top_data,
    int* mask) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    Ftype maxval = -max_dtype<Ftype>();
    int maxidx = -1;
    if (bottom_data_a[index] > bottom_data_b[index]) {
      // only update for very first bottom_data blob (blob_idx == 0)
      if (blob_idx == 0) {
        maxval = bottom_data_a[index];
        top_data[index] = maxval;
        maxidx = blob_idx;
        mask[index] = maxidx;
      }
    } else {
      maxval = bottom_data_b[index];
      top_data[index] = maxval;
      maxidx = blob_idx + 1;
      mask[index] = maxidx;
    }
  }
}

template <typename Ftype, typename Btype>
void EltwiseLayer<Ftype, Btype>::Forward_gpu(const vector<Blob*>& bottom,
    const vector<Blob*>& top) {
  int* mask = nullptr;
  const int count = top[0]->count();
  //  convert to Ftype
  Ftype* top_data = top[0]->mutable_gpu_data<Ftype>();

  switch (op_) {
  case EltwiseParameter_EltwiseOp_PROD:
    caffe_gpu_mul(count, bottom[0]->gpu_data<Ftype>(), bottom[1]->gpu_data<Ftype>(), top_data);
    for (int i = 2; i < bottom.size(); ++i) {
      caffe_gpu_mul(count, top_data, bottom[i]->gpu_data<Ftype>(), top_data);
    }
    break;
  case EltwiseParameter_EltwiseOp_SUM:
    if (shared_ && no_coeffs_) {
      for (int i = 1; i < bottom.size(); ++i) {
        caffe_gpu_incr(count, bottom[i]->gpu_data<Ftype>(), top_data);
      }
    } else {
      if (shared_) {
        caffe_gpu_scal(count, Ftype(coeffs_[0]), top_data);
        for (int i = 1; i < bottom.size(); ++i) {
          caffe_gpu_axpy(count, Ftype(coeffs_[i]), bottom[i]->gpu_data<Ftype>(), top_data);
        }
      } else {
        caffe_gpu_set(count, Ftype(0.), top_data);
        // TODO(shelhamer) does cuBLAS optimize to sum for coeff = 1?
        for (int i = 0; i < bottom.size(); ++i) {
          caffe_gpu_axpy(count, Ftype(coeffs_[i]), bottom[i]->gpu_data<Ftype>(), top_data);
        }
      }
    }
    break;
  case EltwiseParameter_EltwiseOp_MAX:
    {
      mask = max_idx_.mutable_gpu_data();
      cudaStream_t stream = Caffe::thread_stream();
      // NOLINT_NEXT_LINE(whitespace/operators)
      MaxForward <<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, bottom[0]->gpu_data<Ftype>(), bottom[1]->gpu_data<Ftype>(), 0, top_data, mask);
      for (int i = 2; i < bottom.size(); ++i) {
        // NOLINT_NEXT_LINE(whitespace/operators)
        MaxForward <<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
            count, top_data, bottom[i]->gpu_data<Ftype>(), i - 1, top_data, mask);
      }
      CUDA_CHECK(cudaStreamSynchronize(stream));
    }
    break;
  default:
    LOG(FATAL) << "Unknown elementwise operation.";
  }
}

template <typename Btype>
__global__ void MaxBackward(const int nthreads, const Btype* top_diff,
    const int blob_idx, const int* mask, Btype* bottom_diff) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    Btype gradient = 0;
    if (mask[index] == blob_idx) {
      gradient += top_diff[index];
    }
    bottom_diff[index] = gradient;
  }
}

template <typename Ftype, typename Btype>
void EltwiseLayer<Ftype, Btype>::Backward_gpu(const vector<Blob*>& top,
    const vector<bool>& propagate_down, const vector<Blob*>& bottom) {
  const int* mask = nullptr;
  const int count = top[0]->count();
  const Btype* top_data = top[0]->gpu_data<Btype>();
  const Btype* top_diff = top[0]->gpu_diff<Btype>();
  for (int i = 0; i < bottom.size(); ++i) {
    if (propagate_down[i]) {
      const Btype* bottom_data = bottom[i]->gpu_data<Btype>();
      switch (op_) {
      case EltwiseParameter_EltwiseOp_PROD:
        {
          Btype *bottom_diff = bottom[i]->mutable_gpu_diff<Btype>();
          if (stable_prod_grad_) {
            bool initialized = false;
            for (int j = 0; j < bottom.size(); ++j) {
              if (i == j) { continue; }
              if (!initialized) {
                caffe_copy(count, bottom[j]->gpu_data<Btype>(), bottom_diff);
                initialized = true;
              } else {
                caffe_gpu_mul(count, bottom[j]->gpu_data<Btype>(), bottom_diff, bottom_diff);
              }
            }
          } else {
            caffe_gpu_div(count, top_data, bottom_data, bottom_diff);
          }
          caffe_gpu_mul(count, bottom_diff, top_diff, bottom_diff);
        }
        break;
      case EltwiseParameter_EltwiseOp_SUM:
        if (shared_) {
          if (!no_coeffs_) {
            caffe_gpu_scale(count, Btype(coeffs_[i]), top_diff,
                bottom[i]->mutable_gpu_diff<Btype>());
          }
        } else {
          if (no_coeffs_) {
            caffe_copy(count, top_diff, bottom[i]->mutable_gpu_diff<Btype>());
          } else {
            caffe_gpu_scale(count, Btype(coeffs_[i]), top_diff,
                bottom[i]->mutable_gpu_diff<Btype>());
          }
        }
        break;
      case EltwiseParameter_EltwiseOp_MAX:
        mask = max_idx_.gpu_data();
        MaxBackward<Btype>  // NOLINT_NEXT_LINE(whitespace/operators)
            <<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, Caffe::thread_stream()>>>(
            count, top_diff, i, mask, bottom[i]->mutable_gpu_diff<Btype>());
        CUDA_CHECK(cudaStreamSynchronize(Caffe::thread_stream()));
        break;
      default:
        LOG(FATAL) << "Unknown elementwise operation.";
      }
    }
  }
}

  // Implementation of auxiliary boundary-based Forward_gpu/Backward_gpu functions.
  // Added by Hao FU.
  template <typename Ftype, typename Btype>
  void EltwiseLayer<Ftype, Btype>::Forward_gpu(const vector<Blob*>& bottom,
                                               const vector<Blob*>& top,
                                               BatchBoundary_t& boundary) {
    const int bottom_stride_count = bottom[0]->count(1);
    const int top_stride_count = top[0]->count(1);

    int* mask = nullptr;
    const int count = top_stride_count * (boundary.end_idx_ -
                                          boundary.start_idx_ + 1);
    // convert to Ftype.
    Ftype* top_data = top[0]->mutable_gpu_data<Ftype>() +
      boundary.start_idx_ * top_stride_count;

    switch(op_) {
    case EltwiseParameter_EltwiseOp_PROD:
      caffe_gpu_mul(count,
                    bottom[0]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
                    bottom[1]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
                    top_data + boundary.start_idx_ * top_stride_count);
      for (int i = 2; i < bottom.size(); ++ i) {
        caffe_gpu_mul(count, top_data + boundary.start_idx_ * top_stride_count,
                      bottom[i]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
                      top_data + boundary.start_idx_ * top_stride_count);
      }
      break;
    case EltwiseParameter_EltwiseOp_SUM:
      if (shared_ and no_coeffs_) {
        for (int i = 1; i < bottom.size(); ++ i) {
          caffe_gpu_incr(count,
                         bottom[i]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
                         top_data + boundary.start_idx_ * top_stride_count);
        }
      } else {
        if (shared_) {
          caffe_gpu_scal(count,
                         Ftype(coeffs_[0]),
                         top_data + boundary.start_idx_ * top_stride_count);
          for (int i = 1; i < bottom.size(); ++ i) {
            caffe_gpu_axpy(count,
                           Ftype(coeffs_[i]),
                           bottom[i]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
                           top_data + boundary.start_idx_ * top_stride_count);
          }
        } else {
          caffe_gpu_set(count, Ftype(0.),
                        top_data + boundary.start_idx_ * top_stride_count);
          // TODO(shelhamer) does cuBLAS optimize to sum for coeff = 1?
          for (int i = 0; i < bottom.size(); ++ i) {
            caffe_gpu_axpy(count,
                           Ftype(coeffs_[i]),
                           bottom[i]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
                           top_data + boundary.start_idx_ * top_stride_count);
          }
        }
      }
      break;
    case EltwiseParameter_EltwiseOp_MAX:
      {
        mask = max_idx_.mutable_gpu_data() +
          boundary.start_idx_ * bottom_stride_count;
        cudaStream_t stream = Caffe::thread_stream();
        // NOLINT_NEXT_LINE(whitespace/operators)
        MaxForward<<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
            count,
            bottom[0]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
            bottom[1]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
            0, top_data + boundary.start_idx_ * top_stride_count, mask);
        for (int i = 2; i < bottom.size(); ++ i) {
          // NOLINT_NEXT_LINE(whitespace/operators)
          MaxForward<<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
              count,
              top_data + boundary.start_idx_ * top_stride_count,
              bottom[i]->gpu_data<Ftype>() + boundary.start_idx_ * bottom_stride_count,
              i - 1, top_data + boundary.start_idx_ * top_stride_count, mask);
        }
        CUDA_CHECK(cudaStreamSynchronize(stream));
      }
      break;
    default:
      LOG(FATAL) << "Unknown elementwise operation.";
    }
  }

  template <typename Ftype, typename Btype>
  void EltwiseLayer<Ftype, Btype>::Backward_gpu(const vector<Blob*>& top,
                                                const vector<bool>& propagate_down,
                                                const vector<Blob*>& bottom,
                                                BatchBoundary_t& boundary) {
    const int bottom_stride_count = bottom[0]->count(1);
    const int top_stride_count = top[0]->count(1);

    const int* mask = nullptr;
    const int count = top_stride_count *
      (boundary.end_idx_ - boundary.start_idx_ + 1);

    const Btype* top_data = top[0]->gpu_data<Btype>() +
      (boundary.start_idx_ * top_stride_count);
    const Btype* top_diff = top[0]->gpu_diff<Btype>() +
      (boundary.start_idx_ * top_stride_count);
    for (int i = 0; i < bottom.size(); ++ i) {
      if (propagate_down[i]) {
        const Btype* bottom_data = bottom[i]->gpu_data<Btype>() +
          (boundary.start_idx_ * bottom_stride_count);
        switch(op_) {
        case EltwiseParameter_EltwiseOp_PROD:
          {
            Btype* bottom_diff = bottom[i]->mutable_gpu_diff<Btype>() +
              (boundary.start_idx_ * bottom_stride_count);
            if (stable_prod_grad_) {
              bool initialized = false;
              for (int j = 0; j < bottom.size(); ++ j) {
                if (i == j) { continue ; }
                if (!initialized) {
                  caffe_copy(count, bottom[j]->gpu_data<Btype>() + (boundary.start_idx_ * bottom_stride_count),
                             bottom_diff);
                  initialized = true;
                } else {
                  caffe_gpu_mul(count, bottom[j]->gpu_data<Btype>() + (boundary.start_idx_ * bottom_stride_count),
                                bottom_diff, bottom_diff);
                }
              }
            } else {
              caffe_gpu_div(count, top_data, bottom_data, bottom_diff);
            }
            caffe_gpu_mul(count, bottom_diff, top_diff, bottom_diff);
          }
          break;
        case EltwiseParameter_EltwiseOp_SUM:
          if (shared_) {
            if (!no_coeffs_) {
              caffe_copy(count, top_diff,
                         bottom[i]->mutable_gpu_diff<Btype>() + (boundary.start_idx_ * bottom_stride_count));
            } else {
              caffe_gpu_scale(count, Btype(coeffs_[i]), top_diff,
                              bottom[i]->mutable_gpu_diff<Btype>() + (boundary.start_idx_ * bottom_stride_count));
            }
          }
          break;
        case EltwiseParameter_EltwiseOp_MAX:
          mask = max_idx_.gpu_data() + boundary.start_idx_ * bottom_stride_count;
          MaxBackward<Btype>    // NOLINT_NEXT_LINE(whitespace/operators)
            <<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, Caffe::thread_stream()>>>(
               count, top_diff, i, mask,
               bottom[i]->mutable_gpu_diff<Btype>() + boundary.start_idx_ * bottom_stride_count);
          CUDA_CHECK(cudaStreamSynchronize(Caffe::thread_stream()));
          break;
        default:
          LOG(FATAL) << "Unknown elementwise operation.";
        }
      }
    }
  }
  // End. Added by Hao FU.

INSTANTIATE_LAYER_GPU_FUNCS_FB(EltwiseLayer);

// Auxiliary Forward_gpu/Backward_gpu functions for the HGP4NN framework.
// Added by Hao FU.
INSTANTIATE_LAYER_GPU_BOUNDARY_FUNCS_FB(EltwiseLayer);
// End. Added by Hao FU.

}  // namespace caffe
